#!/usr/bin/env python
# Author: Denis Gorbunov
# This package uses FaceL avaialble from http://pyvision.sourceforge.net.
## FaceL official information (required by license):
## Created by David S. Bolme and J. Ross Beveridge
## Colorado State University
## Fort Collins, Colorado, USA'''

import roslib; roslib.load_manifest('facel')
import rospy

from std_msgs.msg import String
from facel.srv import *

from conversation.msg import PersonEvent

from sensor_msgs.msg import Image

import cv
from cv_bridge import CvBridge, CvBridgeError
import opencv

import math
import pickle
import time

import PIL
import PIL.ImageFont
from PIL.Image import FLIP_LEFT_RIGHT

from threading import Thread, Lock, Condition

import pyvision as pv

# The ASEF eye locator has patent complications.  This next line
# disables those warnings.
pv.disableCommercialUseWarnings()

from pyvision.face.CascadeDetector import CascadeDetector
from pyvision.analysis.FaceAnalysis.FaceDetectionTest import is_success
from pyvision.face.SVMFaceRec import SVMFaceRec
from pyvision.face.FilterEyeLocator import loadFilterEyeLocator

CASCADE_NAME = roslib.packages.find_resource('facel', 'haarcascade_frontalface_alt.xml')[0]
FEL_NAME = roslib.packages.find_resource('facel', 'EyeLocatorASEF128x128.fel')[0]
ARIAL_BLACK_NAME = roslib.packages.find_resource('facel', 'Arial Black.ttf')[0]

SVM_MANUAL    = "SVM_MANUAL"
SVM_AUTOMATIC = "SVM_AUTOMATIC"

class Facel():
    def __init__(self):
        rospy.init_node('fasel')
        self.enrollService = rospy.Service("enroll", Enroll, self.handleEnrollFace)
        self.trainService = rospy.Service("train", Train, self.handleTrain)
        self.resetService = rospy.Service("reset", Reset, self.handleReset)
        self.namesPub = rospy.Publisher("names", String)
        self.personEventPub = rospy.Publisher("person_event", PersonEvent)
        self.imagePub = rospy.Publisher("images", Image)
        self.imageSub = rospy.Subscriber("image", Image, self.handleImage)
        
        self.bridge = CvBridge()
        
        self.face_detector = CascadeDetector(cascade_name=CASCADE_NAME,image_scale=0.5)
        self.fel = loadFilterEyeLocator(FEL_NAME)
        self.face_rec = SVMFaceRec()

        self.arialblack24 = PIL.ImageFont.truetype(ARIAL_BLACK_NAME, 24)
        
        self.svm_mode  = SVM_AUTOMATIC
        self.svm_C     = 4.000e+00
        self.svm_Gamma = 9.766e-04
        
        self.current_faces = []
        self.enrolling    = None
        self.enroll_count = 0
        self.enroll_max   = 64
        self.enroll_list  = []
        self.enrollCondition = Condition()

        self.isTrained = False
        
        self.nameList = []
        self.faceNames = []
        self.faceCount = 0
        self.faceTrack = dict()
        
        self.face_processing = True
        self.eye_processing = True
        self.image_flip = True
        
        self.reqQueueLock = Lock()
        self.faceTrackLock = Lock()

        self.is_running = True
        self.eventUpdateDelay = 2.0 # seconds
        self.eventUpdateTimer = Thread(target=self.eventUpdate)
        self.eventUpdateTimer.start()

    # Convert opencv2 image to pyvision image
    def opencv_to_pyvision(self, cv_image):                        
        pil_img = PIL.Image.fromstring("RGB",(cv_image.width,cv_image.height), cv_image.tostring())            
        pyimg = pv.Image(pil_img)                
        return pyimg
    
    # Convert pyvision image to opencv2 image
    def pyvision_to_opencv(self, pyimg):
        pil_img = pyimg.asPIL()
        cv_img = cv.CreateImageHeader(pil_img.size, cv.IPL_DEPTH_8U, 3)  # RGB image
        cv.SetData(cv_img, pil_img.tostring(), pil_img.size[0]*3)
        return cv_img

    def PIL_to_opencv(self, pil_img):
        cv_img = cv.CreateImageHeader(pil_img.size, cv.IPL_DEPTH_8U, 3)  # RGB image
        cv.SetData(cv_img, pil_img.tostring(), pil_img.size[0]*3)
        return cv_img
                                        
    def onFrame(self, img):
        '''
        Process a video frame.
        '''                    
        self.eye_time = 0.0
            
        names = []
        nFaces = 0
        
        if self.face_processing:
        
            faces = self.findFaces(img)
            nFaces = len(faces)
            
            if self.enrolling != None:                
                success = None
                for rect,leye,reye in faces:
                    img.annotateRect(self.enrolling,color='yellow')
                    
                    if (success == None) and is_success(self.enrolling,rect):
                        success = rect
                        img.annotateRect(rect,color='blue')
                        if self.eye_processing:
                            img.annotatePoint(leye,color='blue')
                            img.annotatePoint(reye,color='blue')
                        self.enroll_list.append([img,rect,leye,reye])

                    else:
                        img.annotateRect(rect,color='red')
                        if self.eye_processing:
                            img.annotatePoint(leye,color='red')
                            img.annotatePoint(reye,color='red')
                        img.annotateLine(pv.Point(rect.x,rect.y),pv.Point(rect.x+rect.w,rect.y+rect.h), color='red')
                        img.annotateLine(pv.Point(rect.x+rect.w,rect.y),pv.Point(rect.x,rect.y+rect.h), color='red')

                if success == None:
                    rect = self.enrolling
                    img.annotateLine(pv.Point(rect.x,rect.y),pv.Point(rect.x+rect.w,rect.y+rect.h), color='yellow')
                    img.annotateLine(pv.Point(rect.x+rect.w,rect.y),pv.Point(rect.x,rect.y+rect.h), color='yellow')
                else:
                    #enroll in the identification algorithm
                    pass
            else:
                for rect,leye,reye in faces:
                    img.annotateRect(rect,color='blue')
                    if self.eye_processing:
                        img.annotatePoint(leye,color='blue')
                        img.annotatePoint(reye,color='blue')
                
                for rect,leye,reye in faces:
                    img.annotateRect(rect,color='blue')
                    img.annotatePoint(leye,color='blue')
                    img.annotatePoint(reye,color='blue')
                                   
            if self.isTrained:
                self.label_time = time.time()
                for rect,leye,reye in faces:
                    if (self.face_rec.isTrained()):
                        label = self.face_rec.predict(img,leye,reye)
                        names.append([0.5*(leye+reye),label])
                    
                self.label_time = time.time() - self.label_time                        
            
        im = img.asAnnotated()
            
        # Flip to mirror image
        if self.image_flip:
            im = im.transpose(FLIP_LEFT_RIGHT)
            
        if self.enrolling != None:
            self.enrollCondition.acquire()
            self.enroll_count += 1
            self.enrollCondition.notify()
            self.enrollCondition.release()            
                        
            # Draw on the image
            draw = PIL.ImageDraw.Draw(im)
            x,y = self.enrolling.x,self.enrolling.y
            if self.image_flip:
                xsize, ysize = im.size
                x = xsize - (x + self.enrolling.w)            
            draw.text((x+10,y+10), "Enrolling: %2d of %2d"%(self.enroll_count,self.enroll_max), fill='yellow', font=self.arialblack24)
            del draw
                                                                 
        facesEntered = []

        if len(names) > 0:

            draw = PIL.ImageDraw.Draw(im)

            for pt,name in names:                
                x,y = pt.X(),pt.Y() 
                
                # Draw on the image
                x,y = pt.X(),pt.Y() 
                w,h = draw.textsize(name,font=self.arialblack24)

                if self.image_flip:
                    xsize, ysize = im.size
                    x = xsize - x - 0.5*w
                else:
                    x = x - 0.5*w
                
                draw.text((x,y-20-h), name, fill='green', font=self.arialblack24)

                facesEntered.append(name)
                
                # Publish only new names
                if name not in self.faceNames:
                    str = "seeing %s"%name
                    rospy.loginfo(str)
                    self.namesPub.publish(String(name))
                    self.publishPersonEvent(name, "entered")
                        
            del draw
                    
        # Find all of the faces that are no longer detected        
        for name in self.faceNames:
            if (name not in facesEntered):
                self.publishPersonEvent(name, "exited")
        
      #  print "{0} {1} {2}".format(nFaces, self.faceCount, len(facesEntered))
        nFaces = nFaces - len(facesEntered)

        # For unidentified faces
        # figure out how many entered/exited
        if (nFaces - self.faceCount) > 0:
            self.publishPersonEvent("unknown", "entered")

        if (nFaces - self.faceCount) < 0:
            self.publishPersonEvent("unknown", "exited")
        
        # Update all for the next round
        self.faceNames = facesEntered
        self.faceCount = nFaces;
                
        # Publish the image
        cv_img = self.PIL_to_opencv(im);
        #cv_img = self.pyvision_to_opencv(img)
        msg = self.bridge.cv_to_imgmsg(cv_img, encoding='rgb8')
        
        self.imagePub.publish(msg)

    def publishPersonEvent(self,name,event):
        '''
        Updates the dictionary
        '''
        now = time.time()
                
        # In case of non-existing records, force the update
        if name in self.faceTrack:
            self.faceTrackLock.acquire()
            pastEvent, lastEvent, recTime = self.faceTrack[name]
            self.faceTrackLock.release()
        else:
            recTime = now - self.eventUpdateDelay
            lastEvent = "undefined"
            pastEvent = lastEvent
            
        self.faceTrackLock.acquire()
        self.faceTrack[name] = (pastEvent, event, now)
        self.faceTrackLock.release()
        #print "publishPersonEvent for {0} {1} {2} {3} ".format(name,pastEvent,event,now)

    def eventUpdate(self):
        '''
        Publishes events and updates the dictionary
        '''
        while self.is_running:
            now = time.time()
            # Publish evey n seconds to avoid hysterisis,
            # where n = self.eventUpdateDelay        
            self.faceTrackLock.acquire()
            for name in self.faceTrack:
                lastEvent,event,eventTime = self.faceTrack[name]
                if ((now - eventTime) >= self.eventUpdateDelay) and (event <> lastEvent):
                    #print "eventUpdate for {0} {1} {2} {3} ".format(name,lastEvent,event,now)
                    self.personEventPub.publish(event, name)
                    self.faceTrack[name] = (event, event, now)
                    #print "eventUpdate to {0} ".format(event)
            self.faceTrackLock.release()
                    
            time.sleep(self.eventUpdateDelay)

    def enrollFace(self, name):
        with self.reqQueueLock:
            '''
            Starts the enrollment process.
            '''
            faces = self.current_faces;
            
            if len(faces) > 1:
                return ("too many faces to enroll")        
            
            for rect,leye,reye in faces:
                print "Start Enrolling %s"%(name)
                self.enrolling = rect
                self.enroll_count = 0
                self.nameList.append(name)
                                
                self.enrollCondition.acquire()
                while self.enroll_count < self.enroll_max:
                    self.enrollCondition.wait()
                self.enrollCondition.release()
                
                self.enrolling = None
                
                for data,rect,leye,reye in self.enroll_list:   
                    self.face_rec.addTraining(data,leye,reye,name)      
                                                                  
                self.enroll_count = 0
                self.enroll_list  = []
                    
        return ("ok")
    
    def trainFaces(self):
        with self.reqQueueLock:
            '''
            Start the SVM training process.
            '''
            result = "ok"
            
            # Prevent face prediction while training
            self.isTrained = False

            if self.svm_mode == SVM_AUTOMATIC:
                # Train with automatic tuning.
                print "Training mode: auto tuning"
                self.face_rec.train() #callback=progress.Update)
                self.svm_C = self.face_rec.svm.C
                self.svm_Gamma = self.face_rec.svm.gamma
            else:
                # Train with manual tuning.
                print "Training mode: manual tuning"
                self.face_rec.train(C=[self.svm_C] , Gamma=[self.svm_Gamma])# ,callback=progress.Update)
            
            if (self.face_rec.isTrained()):
                result = "ok"
                self.isTrained = True
            else:
                result = "Error"
                
        return result

    def resetFaces(self):
        with self.reqQueueLock:
            '''
            Clear the enrollment data.
            '''
            result = "ok"

            self.isTrained = False

            nameList = []
            
            # Clear the enrollment data for the SVM
            self.face_rec.reset()
        
        return result
        
    def findFaces(self, img):
        faces = []
        
        self.detect_time = time.time()
        rects = self.face_detector.detect(img)
        self.detect_time = time.time() - self.detect_time

        cvtile = opencv.cvCreateMat(128,128,opencv.CV_8UC3)
        bwtile = opencv.cvCreateMat(128,128,opencv.CV_8U)
        
        cvimg = img.asOpenCV()

        self.eye_time = time.time()
        
        for rect in rects:
            faceim = opencv.cvGetSubRect(cvimg, rect.asOpenCV())
            opencv.cvResize(faceim, cvtile)
            
            affine = pv.AffineFromRect(rect,(128,128))
            
            opencv.cvCvtColor(cvtile, bwtile, cv.CV_BGR2GRAY)                
            
            leye,reye,lcp,rcp = self.fel.locateEyes(bwtile)
            leye = pv.Point(leye)
            reye = pv.Point(reye)
            
            leye = affine.invertPoint(leye)
            reye = affine.invertPoint(reye)
            
            faces.append([rect,leye,reye])
        
        self.eye_time = time.time() - self.eye_time
        self.current_faces = faces
        
        return faces
    
    def handleImage(self, image_message):
        cv_image = self.bridge.imgmsg_to_cv(image_message, desired_encoding='rgb8')
        pyimg = self.opencv_to_pyvision(cv_image)
    
        self.onFrame(pyimg)
    
        #msg = self.bridge.cv_to_imgmsg(cv_image, encoding='rgb8')
        #self.imagePub.publish(msg)

    def handleEnrollFace(self, msg):
        result = "ok"
        print "Enrolling new name:" + msg.name
        
        result = self.enrollFace(msg.name)
        return EnrollResponse(result)                
                        
    def handleTrain(self, msg):        
        result = "ok"
        print "Training..."

        if len(self.nameList) >= 2:
            result = self.trainFaces()
        else:
            result = "Error: need at least 2 faces to train"

        return TrainResponse(result)

    def handleReset(self, msg):
        result = "ok"
        print "Resetting..."
        
        result = self.resetFaces()
            
        return ResetResponse(result)


if __name__ == '__main__':    
    try:
        facel = Facel()
        rospy.spin()
        facel.is_running = False
        
    except rospy.ROSInterruptException:
        pass
